Repository navigation
[mypyc] add support for contextlib.contextmanager - #22267
hauntsaninja wants to merge 7 commits into
Conversation
|
I guess you can say that this fixes mypyc/mypyc#905, right? (Btw I didn't look in details, @JukkaL would be a better reviewer here.) |
| ctm_def_node = None | ||
| # Since we construct the call to the method ourselves in _transform_with_contextmanager, | ||
| # we're a little paranoid about doing things in the wrong order in weird cases | ||
| if not find_walrus_targets(expr) and not expr_has_suspend(expr): |
There was a problem hiding this comment.
would be good to add tests for these cases.
| and len(dec.decorators) == 1 | ||
| and isinstance(dec1 := dec.decorators[0], mypy.nodes.RefExpr) |
There was a problem hiding this comment.
also add a test where the function has other decorators, i don't think we currently have one like that.
| wrapped_call = mypy.nodes.CallExpr( | ||
| mypy.nodes.MemberExpr(expr.callee, "__wrapped__"), | ||
| expr.args, | ||
| expr.arg_kinds, | ||
| expr.arg_names, | ||
| ) | ||
| wrapped_call.line = line | ||
| gen = builder.accept(wrapped_call) |
There was a problem hiding this comment.
i think this results in a generic Python attribute access + generic function call at runtime which we could optimize for compiled functions decorated with @contextmanager. for those we could call the generated C function directly.
| assert isinstance(expr.callee, mypy.nodes.MemberExpr) | ||
| receiver = builder.accept(receiver_expr) | ||
| method = builder.py_get_attr(receiver, expr.callee.name, line) | ||
| wrapped = builder.py_get_attr(method, "__wrapped__", line) |
There was a problem hiding this comment.
for non-final methods we should check and optimize only if __wrapped__ actually exists because it's legal to override a contextmanager method with a regular method in a subclass.
we might be able to optimize the call for final methods and for methods of compiled classes that don't allow interpreted subclasses by calling the C function directly.
| def test_no_yield() -> None: | ||
| try: | ||
| with cm_no_yield(): | ||
| print("UNREACHABLE") | ||
| print("UNREACHABLE") | ||
| except RuntimeError as e: | ||
| e_any = cast(Any, e) | ||
| print( | ||
| type(e).__name__, | ||
| str(e), | ||
| e_any.__cause__ is None, | ||
| e_any.__suppress_context__, | ||
| ) |
There was a problem hiding this comment.
general comment about the run tests: when adding new tests we usually check using inline assertions instead of checking the printed output. i would prefer to use inline assertions here as well since it makes the tests easier to read - i don't have to scroll down to see the desired behavior.
| try: | ||
| with cm_no_stop() as v: | ||
| print("entered body", v) | ||
| print("UNREACHABLE") | ||
| except RuntimeError as e: | ||
| print(type(e).__name__, str(e)) |
There was a problem hiding this comment.
general comment: could use assertRaises for checking exceptions.
|
|
||
| print("ret result", inner()) | ||
|
|
||
| [file driver.py] |
There was a problem hiding this comment.
if you use inline assertions i think you could also simplify the new tests by not using a custom driver.py. if there's no custom one then the default runs every function in the compiled file prefixed with test_.
| with cm_no_yield(): | ||
| print("UNREACHABLE") |
There was a problem hiding this comment.
would be good to also add an irbuild test with a specialized with statement since the run tests don't prove that the statement was compiled using the specialization.
|
Thanks for the review, will address! I used prints instead of runtime asserts because the exception logic is actually kind of tricky and was a source of bugs for me, very easy to accidentally suppress an exception I couldn't see an easy way to find the generated C function. Codex had a diff where it did something that created some new statics, it benchmarked only a little bit better, but you're right that it could be a little more robust. Let me investigate it. |
Should make self check like 6-10% faster and torch 8-12% faster (compiled, clean cache, not parallel)
Fixes mypyc/mypyc#905
Co-authored by Codex